from llama_cpp import Llama

# 初始化模型
# llm = Llama(
#     model_path="./llama.cpp/deepseek-r1-1.5b.q4_0.gguf",  # GGUF模型路径
#     n_ctx=2048,          # 上下文窗口大小（根据模型调整）
#     n_threads=8,         # CPU线程数
#     n_gpu_layers=35      # 使用GPU加速的层数（如支持，Metal/CUDA）
# )

# 生成文本
# prompt = "法国的首都是哪里?"
# output = llm(
#     prompt=prompt,
#     max_tokens=2048,      # 生成的最大token数
#     stop=["\n"],        # 停止条件（可选）
#     temperature=0.8,    # 控制随机性（0-1，越大越随机）
#     echo=True           # 是否返回输入的prompt
# )
#
# print(output["choices"][0]["text"])


# response = llm.create_chat_completion(
#     messages=[
#         {"role": "system", "content": "你是一个得力的AI助手"},
#         {"role": "user", "content": "如何学好python?"}
#     ],
#     temperature=0.7,
#     max_tokens=2048
# )
#
# print(response["choices"][0]["message"]["content"])


from llama_cpp import Llama

# 初始化模型
llm = Llama(
    model_path="./llama.cpp/deepseek-r1-1.5b.q4_k_m.gguf",
    n_ctx=4096,
    n_gpu_layers=40,
    n_threads=8
)

# 生成回复
prompt = "Translate 'Hello, world!' to French."
response = llm(
    prompt=prompt,
    max_tokens=100,
    temperature=0.7,
    top_p=0.9
)

print(response["choices"][0]["text"])
